package com.bjfu.intelligentrecognizeserver.question_answer;

import com.alibaba.fastjson.JSONObject;
import com.bjfu.intelligentrecognizeserver.conversation.ConversationContext;
import com.bjfu.intelligentrecognizeserver.pojo.entity.UploadImageInfo;
import com.bjfu.intelligentrecognizeserver.pojo.vo.SpeciesVo;
import com.bjfu.intelligentrecognizeserver.service.SpeciesService;
import com.bjfu.intelligentrecognizeserver.service.UploadService;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonPrimitive;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.*;


@Component

public class IntelligentRecognitionServer implements WebSocketHandler {
    private static final Logger log = LoggerFactory.getLogger(IntelligentRecognitionServer.class);
    private WebSocketSession mSession;
    private String userId;
    private final UploadService uploadService;
    private final SpeciesService speciesService;
    //ip
    private static final String IP = "172.27.135.152";
    //172.25.100.207
    private static final String PORT = "10000";
    private static final String COMMUNICATE_WITH_BIG_MODEL_URL = "http://" + IP + ":" + PORT + "/CommunicateWithBigModel";
    private static final String MULTI_MODAL_RECOGNIZE_URL = "http://" + IP + ":" + PORT + "/MultiModalRecognize/";

    @Autowired
    public IntelligentRecognitionServer(UploadService uploadService, SpeciesService speciesService) {
        this.uploadService = uploadService;
        this.speciesService = speciesService;
    }

    private static final ExecutorService threadPool = Executors.newCachedThreadPool();

    // 存储每个会话的上下文
    private static final Map<String, ConversationContext> conversationMap = new HashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("websocket connection established");
        String sessionId = session.getId();
        log.info("session id:{} ", sessionId);
        // 创建一个新的对话上下文
        conversationMap.put(sessionId, new ConversationContext(sessionId));
        String uri = Objects.requireNonNull(session.getUri()).toString();
        log.info("uri:{} ", uri);
        this.userId = uri.substring(uri.lastIndexOf("/") + 1);
        log.info("userId:{} ", this.userId);
        log.info("websocket connection established for userId: {}", userId);
        this.mSession = session;
    }

    @Override
    public void handleMessage(@NotNull WebSocketSession session, @NotNull WebSocketMessage<?> message) throws Exception {
        //此处应该把message传给大模型，大模型回答得到答案，返回给服务器，服务器再将答案给客户端
        String clientMessage = (String) message.getPayload();
        log.info("received message:{}", clientMessage);
        String sessionId = session.getId();
        ConversationContext context = conversationMap.get(sessionId);


        // 使用 Gson 解析 JSON 字符串
        Gson gson = new Gson();
        JsonObject jsonObject = gson.fromJson(clientMessage, JsonObject.class);

        // 获取 text 字段的值
        String userInput = jsonObject.get("text").getAsString();
        CompletableFuture<ResponseData> future = processUserInput(userInput);
        ResponseData responseData;
        try {
            responseData = future.get(100, TimeUnit.SECONDS); // 设置100秒超时
            log.info("answer:{}", responseData.getContent());
        } catch (Exception e) {
            responseData = new ResponseData("error", "请求处理超时，请稍后重试");
            log.error("处理超时: {}", e.getMessage());
        }

        // 修改2：增强空值保护
        if (responseData.getContent() == null) {
            responseData = new ResponseData("error", "服务暂时不可用，请稍后重试");
        }

        context.addMessage(clientMessage); // 保存用户问题
        context.addMessage(responseData.getContent()); // 保存模型的回答
        sseResponse(responseData);

    }

    private void sseResponse(ResponseData responseData) throws IOException {
        String type = responseData.getType();
        String content = responseData.getContent();
        SpeciesVo speciesVo = responseData.getSpeciesVo();
        int chunkSize = 5;  // 每次发送100个字符
        int index = 0;
        mSession.sendMessage(new TextMessage("Response: "));
        // 每100毫秒分块发送一次数据
        while (index < content.length()) {
            String chunk = content.substring(index, Math.min(index + chunkSize, content.length()));
            String prefixedChunk = addTypePrefix(type, chunk);
            mSession.sendMessage(new TextMessage(prefixedChunk));
            index += chunkSize;

            try {
                Thread.sleep(100);  // 控制发送频率
            } catch (InterruptedException e) {
                log.error(e.getMessage());
            }
        }
        mSession.sendMessage(new TextMessage("Response end"));
        // 如果是多模态识别结果且有物种信息，发送物种信息
        if ("multimodal".equals(type) && speciesVo != null) {
            sendSpeciesInfo(speciesVo);
        }
    }

    // 新增方法：发送物种信息
    private void sendSpeciesInfo(SpeciesVo speciesVo) throws IOException {
        Gson gson = new Gson();
        String speciesInfoJson = gson.toJson(speciesVo);
        mSession.sendMessage(new TextMessage("SPECIES_INFO:" + speciesInfoJson));
    }

        // 新增方法：添加类型前缀
        private String addTypePrefix(String type, String chunk) {
            if ("multimodal".equals(type)) {
                return "[多模态] " + chunk;
            } else if ("conversation".equals(type)) {
                return "[对话] " + chunk;
            } else {
                return "[错误] " + chunk;
            }
        }


    //大模型对话
    private CompletableFuture<String> communicateToBigModel(String question) {
        log.info("question:{}", question);
        return CompletableFuture.supplyAsync(() -> {
            HttpURLConnection connection = null;
            BufferedReader in = null;
            String resultString;

            try {
                URL url = new URL(COMMUNICATE_WITH_BIG_MODEL_URL);
                connection = (HttpURLConnection) url.openConnection();
                connection.setRequestMethod("POST");
                connection.setRequestProperty("Content-Type", "application/json");
                connection.setDoOutput(true);

                JsonObject json = new JsonObject();
                json.add("content", new JsonPrimitive(question));
                String jsonInputString = json.toString();

                try (OutputStream os = connection.getOutputStream()) {
                    byte[] input = jsonInputString.getBytes(StandardCharsets.UTF_8);
                    os.write(input, 0, input.length);
                }

                int responseCode = connection.getResponseCode();

                // 修改3：处理非200响应
                if (responseCode == HttpURLConnection.HTTP_OK) {
                    in = new BufferedReader(new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8));
                } else {
                    in = new BufferedReader(new InputStreamReader(connection.getErrorStream(), StandardCharsets.UTF_8));
                }

                StringBuilder response = new StringBuilder();
                String inputLine;
                while ((inputLine = in.readLine()) != null) {
                    response.append(inputLine);
                }

                // 修改4：增强错误处理
                if (responseCode == HttpURLConnection.HTTP_OK) {
                    resultString = response.toString().replace("\\n", "").replace("\"", "");
                } else {
                    resultString = "请求失败，状态码：" + responseCode;
                }
            } catch (Exception e) {
                resultString = "处理异常: " + e.getMessage();
                log.error("处理异常: {}", e.getMessage(), e);
            } finally {
                try {
                    if (in != null) in.close();
                    if (connection != null) connection.disconnect();
                } catch (IOException ex) {
                    log.error("关闭连接异常: {}", ex.getMessage(), ex);
                }
            }
            return resultString;
        }, threadPool);
    }

    //多模态识别
    private CompletableFuture<String> MultiModalRecognize(String userInput) {
        return CompletableFuture.supplyAsync(() -> {
            String result = "";
            try {

                // Path to the image file
                List<UploadImageInfo> infos = uploadService.getUploadImageInfo(userId);
                if (infos.isEmpty()) {
                    return communicate(userInput);
                }
                UploadImageInfo lastInfo = infos.get(infos.size() - 1);
                if (lastInfo.getState() == 0) {
                    return communicate(userInput);
                }
                String nowImagePath = lastInfo.getImagePath();
                log.info("imagePath:{}", nowImagePath);

                uploadService.alterStateById(lastInfo.getId());

                JSONObject json = new JSONObject();
                json.put("image_path", nowImagePath);
                json.put("user_input", userInput);


                String response = sendPostRequest(json.toString());


                Gson gson = new Gson();
                JSONObject jsonObject = gson.fromJson(response, JSONObject.class);
                result = ((String) jsonObject.get("result")).strip();


            } catch (Exception e) {
                log.error(e.getMessage());
            }
            return result;
        }, threadPool);
    }

    // 新增方法：处理用户输入，返回包含类型和内容的响应
    private CompletableFuture<ResponseData> processUserInput(String userInput) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                List<UploadImageInfo> infos = uploadService.getUploadImageInfo(userId);
                if (infos.isEmpty() || infos.get(infos.size() - 1).getState() == 0) {
                    // 对话处理
                    String conversationResult = communicate(userInput);
                    return new ResponseData("conversation", conversationResult, null);
                } else {
                    // 多模态识别
                    UploadImageInfo lastInfo = infos.get(infos.size() - 1);
                    String nowImagePath = lastInfo.getImagePath();
                    uploadService.alterStateById(lastInfo.getId());

                    JSONObject json = new JSONObject();
                    json.put("image_path", nowImagePath);
                    json.put("user_input", userInput);
                    String response = sendPostRequest(json.toString());

                    Gson gson = new Gson();
                    JSONObject jsonObject = gson.fromJson(response, JSONObject.class);
                    String speciesName = ((String) jsonObject.get("result")).strip();


                    // 获取物种详细信息
                    SpeciesVo speciesInfo = speciesService.getSpeciesInfoByName(speciesName);
                    return new ResponseData("multimodal", speciesName, speciesInfo);

                }
            } catch (Exception e) {
                log.error(e.getMessage());
                return new ResponseData("error", "处理过程中发生错误: " + e.getMessage(), null);
            }
        }, threadPool);
    }

    private String communicate(String userInput) {
        CompletableFuture<String> future = communicateToBigModel(userInput);
        String answer;
        try {
            answer = future.get(100, TimeUnit.SECONDS); // 设置100秒超时
        } catch (Exception e) {
            answer = "请求处理超时，请稍后重试";
        }
        return answer;
    }

    private String sendPostRequest(String jsonPayload) throws IOException, InterruptedException {
        HttpClient client = HttpClient.newHttpClient();
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(IntelligentRecognitionServer.MULTI_MODAL_RECOGNIZE_URL))
                .header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(jsonPayload))
                .build();
        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        return response.body();
    }

    //实现真实的流式，但是有一点问题
    private void processQuestionAsyncSse(String question, WebSocketSession session) {
        try {
            // WebSocket 连接地址（确保与服务器地址一致）
            String uri = "ws://172.27.130.255:9000/ws/process";
            URI serverUri = new URI(uri);

            CountDownLatch latch = new CountDownLatch(1);

            WebSocketClient client = new WebSocketClient(serverUri) {
                @Override
                public void onOpen(ServerHandshake handshake) {
                    System.out.println("已连接到服务器");
                    // 发送JSON格式消息
                    //String jsonMessage = "{\"content\":\"马缨丹是什么\"}";
                    JsonObject json = new JsonObject();
                    json.add("content", new JsonPrimitive(question));
                    this.send(json.toString());
                }

                @Override
                public void onMessage(String message) {
                    try {
                        session.sendMessage(new TextMessage(message));
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }

                }

                @Override
                public void onClose(int code, String reason, boolean remote) {
                    System.out.println("连接关闭: " + reason);
                    latch.countDown();
                }

                @Override
                public void onError(Exception ex) {
                    log.error(ex.getMessage());
                    latch.countDown();
                }
            };

            client.connect();  // 建立连接

            // 等待连接关闭
            latch.await();
        } catch (URISyntaxException | InterruptedException e) {
            log.error(e.getMessage());
        }
    }


    @Override
    public void handleTransportError(@NotNull WebSocketSession session, @NotNull Throwable exception) {

    }

    @Override
    public void afterConnectionClosed(@NotNull WebSocketSession session, @NotNull CloseStatus closeStatus) {

    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    //响应数据类
    @Data
    @Builder
    @AllArgsConstructor
    private static class ResponseData{
        private  String type;
        private  String content;
        private  SpeciesVo speciesVo;

        public ResponseData(String type, String content){
            this.type = type;
            this.content = content;
        }
    }
}
